
#include <openrave-core.h>
#include <vector>
#include <iostream>

#include "orexample.h"
#include "pidcontroller.hpp"
#include "jointspacesolver.hpp"
#include "trajectory.hpp"
#include "spline.hpp"
#include "robotStruct.hpp"

#define TIMESTEP 10 //milliseconds

#define ZOFFSET 1.10

#define LINKLEN1 0.0
#define LINKLEN2 0.100
#define LINKLEN3 0.600
#define LINKLEN4 0.500
#define LINKLEN5 0.500

using std::cout;
using std::endl;
using std::string;

namespace cppexamples {

class RobotArm : public OpenRAVEExample {
 public:
  virtual void demothread(int argc, char ** argv) {
    
    // load the world:
    if( !penv->Load("models/world.xml") ){
      std::cout << "Failed to load xml\n";
      return;
    }
    // attach a physics engine:
    penv->SetPhysicsEngine(RaveCreatePhysicsEngine(penv, "ode"));
    penv->GetPhysicsEngine()->SetGravity(OpenRAVE::Vector(0,0,-9.8));
    
    // Find the first robot and get a pointer to it:
    std::vector<OpenRAVE::RobotBasePtr> orRobotVec;
    penv->GetRobots(orRobotVec);
    if(orRobotVec.size() < 1) {
      std::cout << "No robots found in xml\n";
      return;
    }
    
    OpenRAVE::RobotBasePtr orRobot = orRobotVec.at(0);
    uint32_t numJoints = orRobot->GetJoints().size();
    
    std::cout << "Init Robot Model:\n";
    RobotModel* robot = new RobotModel;
    robot->setLengths(LINKLEN1, LINKLEN2, LINKLEN3, LINKLEN4, LINKLEN5);
    
    cout << "Forming Joint Trajectory...\n";
    assert(numJoints == 4);
    
    // ---------------- PLAN PATH ---------------- //
    
    JointSpaceSolver* solver = new JointSpaceSolver(robot);
    
    double reach = solver->getReachableLimit()-0.0001;
    
    Trajectory* trajectory = new Trajectory();
    Setpoint* sp = new Setpoint(0, true);
    
    sp->time = 0;
    sp->cpos->set(reach, 0.0, 0.0, M_PI/3);
    if(solver->solve(sp) != 0) return;
    trajectory->addSetpoint(sp);
    
    
    // ---- WOBBLY CIRCLE ----
    double path_r = reach/2;
    double path_a = reach/6;
    
    sp->time = 1000;
    sp->cpos->set(path_r*cos(0), path_r*sin(0), 0, M_PI/2);
    if(solver->solve(sp) != 0) return;
    trajectory->addSetpoint(sp);
    
    for(double angle = 0; angle < 2*M_PI; angle += M_PI/180){
      sp->time = sp->time + 100;
      double r = path_r + path_a*sin(angle*6);
      sp->cpos->set(r*cos(angle), r*sin(angle), 0, M_PI/2-M_PI*sin(angle*6)/12);
      if(solver->solve(sp) != 0) return;
      trajectory->addSetpoint(sp);
    }
    /*
    // ---- STRAIGHT LINE ----
    double path_x = 3*reach/4;
    double path_y = reach/2;
    
    sp->time = 3000;
    sp->cpos->set(path_x, path_y, 0, M_PI/3);
    if(solver->solve(sp) != 0) return;
    trajectory->addSetpoint(sp);
    
    for(path_x = 3*reach/4; path_x > -3*reach/4; path_x -= 0.05){
      sp->time = sp->time+500;
      sp->cpos->set(path_x, path_y, 0, M_PI/3);
      if(solver->solve(sp) != 0) return;
      trajectory->addSetpoint(sp);
    }
    */
    // Construct Controllers:
    cout << "Constructing Controllers...\n";
    
    std::vector<PIDController*> controllers;
    controllers.resize(numJoints);
    
    // PID Gains:
    double Kp = 200.0;
    double Ki = 0.7;
    double Kd = 1000;
    
    // Max Joint Strength:
    double maxTorque = 100; // Nm
    
    // Joint 0:
    controllers[0] = new PIDController(Kp, Ki, Kd);
    controllers[0]->setAntiWindup(-maxTorque, maxTorque);

    // Joint 1:
    controllers[1] = new PIDController(Kp, Ki, Kd);
    controllers[1]->setAntiWindup(-maxTorque, maxTorque);
    
    // Joint 2:
    controllers[2] = new PIDController(Kp, Ki, Kd);
    controllers[2]->setAntiWindup(-maxTorque, maxTorque);
    
    // Joint 3:
    controllers[3] = new PIDController(Kp, Ki, Kd);
    controllers[3]->setAntiWindup(-maxTorque, maxTorque);

    cout << "Entering loop...\n";
    uint64_t simtime = 0.0;
    
    std::vector<OpenRAVE::dReal> positions;        // angular position
    positions.resize(numJoints);
    
    std::vector<OpenRAVE::dReal> torques;          // joint torques
    torques.resize(numJoints);
    
    //bool drawLines = true;
    
    while(IsOk()) {
      {  // Begin Mutually Exclusive Segment:
        OpenRAVE::EnvironmentMutex::scoped_lock lock(penv->GetMutex()); // lock environment
        
        // Control:
        simtime = penv->GetSimulationTime()/1000;    // milliseconds
        //cout << "time = " << simtime << "ms \n";
        // Get Joint positions:
        orRobot->GetDOFValues(positions);
        
        // Get Joint set points:
        trajectory->getSetpoint(simtime, sp);   // to repeat use (simtime%12000)
        
        for(uint32_t i=0; i<numJoints; i++) {
          /*
          cout << "Joint:" << i << "  Pos:" << positions.at(i) 
               << "  Set Point:" << sp->at(i) << "\n";
          */
          // Get control signal from PID:
          double controlSignal = controllers.at(i)->compute(positions.at(i), 
                                                         sp->jpos->get(i));
          
          // Max/Min to Actuator rails:
          double torque = controlSignal;
          if( controlSignal >  maxTorque ) torque =  maxTorque;
          if( controlSignal < -maxTorque ) torque = -maxTorque;
          
          torques.at(i) = torque;
          //cout << "  Torque:" << torques.at(i) << "\n";
        }
        // Set joint torques
        orRobot->SetDOFTorques(torques, false);
      }  // end Mutually Exclusive Segment
      
      // Drawing:
      std::list<OpenRAVE::GraphHandlePtr> listgraphs;
      listgraphs.clear();
      std::vector< OpenRAVE::RaveVector<float> > vpoints;
      
      Setpoint* sp1 = new Setpoint(0, true);
      
      for(long ti = 0; ti < 50000; ti += 500) {
        vpoints.clear();
        
        trajectory->getSetpoint(ti, sp1);
        OpenRAVE::RaveVector<float> rv1(sp1->cpos->x, sp1->cpos->y, sp1->cpos->z+ZOFFSET);
        vpoints.push_back(rv1);
        
        trajectory->getSetpoint(ti+500, sp1);
        //sp->println();
        OpenRAVE::RaveVector<float> rv2(sp1->cpos->x, sp1->cpos->y, sp1->cpos->z+ZOFFSET);
        vpoints.push_back(rv2);
        
        listgraphs.push_back(penv->drawlinestrip(&vpoints[0].x, vpoints.size(), sizeof(vpoints[0]), 1.0f));
      }
      vpoints.clear();
      
      OpenRAVE::RaveVector<float> rv3(0, 0, 0);
      vpoints.push_back(rv3);
      
      cvec4d* cpos = new cvec4d();
      solver->solveFK(sp->jpos, cpos);
      
      OpenRAVE::RaveVector<float> rv4(cpos->x, cpos->y, cpos->z+ZOFFSET);
      vpoints.push_back(rv4);
      
      listgraphs.push_back(penv->drawlinestrip(&vpoints[0].x, vpoints.size(), sizeof(vpoints[0]), 1.0f));
      
      vpoints.clear();
      
      OpenRAVE::RaveVector<float> color(255, 255, 255, 255);
      
      OpenRAVE::RaveVector<float> rv5(0, 0, 0);
      vpoints.push_back(rv5);
      
      jvec4d* jpos = new jvec4d(positions[0], positions[1], positions[2], positions[3]);
      
      solver->solveFK(jpos, cpos);
      
      OpenRAVE::RaveVector<float> rv6(cpos->x, cpos->y, cpos->z+ZOFFSET);
      vpoints.push_back(rv6);
      
      listgraphs.push_back(penv->drawlinestrip(&vpoints[0].x, vpoints.size(), sizeof(vpoints[0]), 1.0f, color));
      
      delete cpos;
      delete jpos;
      
      while( penv->GetSimulationTime()/1000 - simtime < TIMESTEP && IsOk()) {
        // wait...
      }
    }
    // Cleanup:
    delete sp;
    delete robot;
    delete solver;
    delete trajectory;
    for(uint32_t i = 0; i < controllers.size(); i ++) {
      delete controllers.at(i);
    }
  }
};  // end class

} // end namespace cppexamples
